"""
Copyright (c) Cookie Yang. All right reserved.
"""
#!/usr/bin/python

from bcc import BPF

# Define BPF program
bpf_txt = """

#include <linux/sched.h>

// Default key value is u64, but PID is u32
BPF_HASH(start, u32);

BPF_PERF_OUTPUT(events);

struct data_t {
    u32 pid;
    u64 ts;
    char comm[TASK_COMM_LEN];
    u64 exe_time;
};

int start_execve(struct pt_regs *ctx) {
    u64 ts;
    u32 pid;

    // Get PID and timestamp (in ns)
    pid = bpf_get_current_pid_tgid();
    ts = bpf_ktime_get_ns();

    // Save the start time (with key=PID) and return
    start.update(&pid, &ts);
    return 0;
}

int end_execve(struct pt_regs *ctx) {
    u64 *tsp, ts, delta;
    struct data_t data = {};

    // Get PID and end timestamp (in ns)
    data.pid = bpf_get_current_pid_tgid();
    ts = bpf_ktime_get_ns();

    // Get the process's name
    bpf_get_current_comm(&data.comm, sizeof(data.comm));

    // Lookup start time for corresponding PID
    tsp = start.lookup(&data.pid);
    if (tsp != NULL) {

        // Start time is also the timestamp, find delta
        data.ts = *tsp;
        data.exe_time = ts - *tsp;

        // Delete the hash entry and submit data
        start.delete(&data.pid);
        events.perf_submit(ctx, &data, sizeof(data));
    }

    return 0;
}
"""

def handle_event(cpu, data, size):
    output = bpf_ctx["events"].event(data)
    print("%-18d %-16s %-6d %-16d" %
          (output.ts, output.comm, output.pid, output.exe_time))

# Load BPF program
bpf_ctx = BPF(text=bpf_txt)
event_name = bpf_ctx.get_syscall_fnname("execve")
bpf_ctx.attach_kprobe(event="event_name", fn_name="start_execve")
bpf_ctx.attach_kretprobe(event="event_name", fn_name="end_execve")

# Print header
print("%-18s %-16s %-6s %-16s" %
      ("TIME(ns)", "COMM", "PID", "EXE TIME(ns)"))

# Open perf "events" buffer
bpf_ctx["events"].open_perf_buffer(handle_event)

# Poll for events
while 1:
    try:
        bpf_ctx.perf_buffer_poll()
    except KeyboardInterrupt:
        print()
        exit()